# voxelizer_open3d.py

"""
使用 Open3D 库进行体素化的高性能引擎模块。
(V11 - 移除硬编码版)
"""
import open3d as o3d
import open3d.core as o3c
import numpy as np
from typing import List, Dict, Any, Tuple
from tqdm import tqdm
from .constants import DEFAULT_MATERIAL_NAME

def legacy_mesh_to_t_mesh(mesh_legacy, device):
    vertices_t = o3c.Tensor(np.asarray(mesh_legacy.vertices), dtype=o3c.float32, device=device)
    triangles_t = o3c.Tensor(np.asarray(mesh_legacy.triangles), dtype=o3c.int32, device=device)
    t_mesh = o3d.t.geometry.TriangleMesh(device)
    t_mesh.vertex['positions'] = vertices_t
    t_mesh.triangle['indices'] = triangles_t
    return t_mesh

def voxelize_with_open3d(model_data: Dict[str, Any], voxel_size: float, solid: bool = True) -> Tuple[List[Dict[str, Any]], Tuple[float, float, float], Tuple[int, int, int]] | Tuple[None, None, None]:
    mode = "实心 (Raycasting)" if solid else "表面 (Surface)"
    print(f"--- 使用 Open3D 引擎 ({mode}) ---")
    try:
        if not model_data or not model_data.get("vertices"):
            print("  错误：传入的模型数据无效。")
            return None, None, None
        print("  步骤 1: 根据传入数据构建几何体和材质映射...")
        vertices_list = model_data["vertices"]
        mesh_legacy = o3d.geometry.TriangleMesh()
        mesh_legacy.vertices = o3d.utility.Vector3dVector(vertices_list)
        triangle_to_material_map: Dict[int, str] = {}
        triangles_list_flat = []
        current_triangle_index = 0
        for grp in model_data["face_groups"]:
            material_name = grp.get("material") or DEFAULT_MATERIAL_NAME
            for face in grp["faces"]:
                if len(face) >= 3:
                    for i in range(1, len(face) - 1):
                        triangle = [face[0], face[i], face[i+1]]
                        triangles_list_flat.append(triangle)
                        triangle_to_material_map[current_triangle_index] = material_name
                        current_triangle_index += 1
        mesh_legacy.triangles = o3d.utility.Vector3iVector(triangles_list_flat)
        if not mesh_legacy.has_triangles():
            return [], (0,0,0), (1,1,1)
        print("  步骤 2: 对整个模型进行体素化...")
        device = o3c.Device("CUDA:0") if hasattr(o3d, 'cuda') and o3d.cuda.is_available() else o3c.Device("CPU:0")
        if solid:
            bbox = mesh_legacy.get_axis_aligned_bounding_box()
            bbox_min, bbox_max = bbox.min_bound - voxel_size, bbox.max_bound + voxel_size
            dims = (bbox_max - bbox_min) / voxel_size
            grid_dims_np = np.ceil(dims).astype(int)
            grid_dims_np = np.maximum(grid_dims_np, 1)
            x, y, z = (np.arange(grid_dims_np[i]) * voxel_size + bbox_min[i] + voxel_size / 2 for i in range(3))
            grid_x, grid_y, grid_z = np.meshgrid(x, y, z, indexing='ij')
            query_points = np.stack([grid_x.ravel(), grid_y.ravel(), grid_z.ravel()], axis=-1)
            t_mesh = legacy_mesh_to_t_mesh(mesh_legacy, device)
            scene = o3d.t.geometry.RaycastingScene()
            scene.add_triangles(t_mesh)
            occupancy = scene.compute_occupancy(o3c.Tensor(query_points, dtype=o3c.float32, device=device))
            inside_indices_np = (occupancy.to(o3c.int32) == 1).cpu().numpy()
            grid_indices = np.stack(np.meshgrid(np.arange(grid_dims_np[0]), np.arange(grid_dims_np[1]), np.arange(grid_dims_np[2]), indexing='ij'), axis=-1)
            active_voxel_indices = grid_indices.reshape(-1, 3)[inside_indices_np]
            grid_origin = tuple(bbox_min)
            grid_dimensions = tuple(grid_dims_np)
        else:
            voxel_grid = o3d.geometry.VoxelGrid.create_from_triangle_mesh(mesh_legacy, voxel_size)
            active_voxel_indices = np.array([v.grid_index for v in voxel_grid.get_voxels()])
            grid_origin = tuple(voxel_grid.origin)
            max_indices = np.max(active_voxel_indices, axis=0) if active_voxel_indices.any() else np.array([0,0,0])
            grid_dimensions = tuple(int(c+1) for c in max_indices)
        if len(active_voxel_indices) == 0:
            return [], grid_origin, grid_dimensions
        print(f"  步骤 3: 为 {len(active_voxel_indices)} 个体素分配材质...")
        voxel_centers_world = active_voxel_indices * voxel_size + np.array(grid_origin) + voxel_size / 2
        t_mesh_coloring = legacy_mesh_to_t_mesh(mesh_legacy, device)
        scene_coloring = o3d.t.geometry.RaycastingScene()
        scene_coloring.add_triangles(t_mesh_coloring)
        closest_points_data = scene_coloring.compute_closest_points(o3c.Tensor(voxel_centers_world, dtype=o3c.float32, device=device))
        closest_triangle_ids = closest_points_data['primitive_ids'].cpu().numpy()
        active_voxels_list = []
        progress_bar_desc = "  └ 分配材质"
        for i, voxel_idx in enumerate(tqdm(active_voxel_indices, desc=progress_bar_desc, ncols=80)):
            triangle_id = closest_triangle_ids[i]
            material_name = triangle_to_material_map.get(triangle_id, DEFAULT_MATERIAL_NAME)
            active_voxels_list.append({
                "x": int(voxel_idx[0]), "y": int(voxel_idx[1]), "z": int(voxel_idx[2]),
                "material": material_name
            })
        print(f"    - 材质分配完成。")
        return active_voxels_list, grid_origin, grid_dimensions
    except Exception as e:
        print(f"错误 (Open3D 引擎): 在处理过程中发生未知异常: {e}")
        import traceback
        traceback.print_exc()
        return None, None, None